package socket

import (
	"bytes"
	comJwt "common/jwt"
	"common/response"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/gorilla/websocket"
	"github.com/rs/xid"
	"io"
	"log"
	"main/svc"
	"main/types"
	"net/http"
	"sync"
	"time"
)

var (
	// AES 和 HMAC 密钥
	secretKey = generateAESKey() // AES 密钥
	hmacKey   = generateAESKey() // HMAC 密钥
)

func generateAESKey() []byte {
	s := []byte("your-aes-16byte-key-1234")
	return s
}

type WebSocketMessage struct {
	EncryptedMessage string `json:"encryptedMessage"`
	Hmac             string `json:"hmac"`
}
type Client struct {
	ID            string          // 连接ID
	AccountId     string          // 账号id, 一个账号可能有多个连接
	Socket        *websocket.Conn // 连接
	HeartbeatTime int64           // 前一次心跳时间
}

func (c *Client) Read() {
	defer func() {
		_ = c.Socket.Close()
	}()
	for {
		// 读取消息
		_, body, err := c.Socket.ReadMessage()
		if err != nil {
			break
		}
		var msg1 WebSocketMessage
		if err := json.Unmarshal(body, &msg1); err != nil {
			log.Println("JSON 解析失败:", err)
			sendErr(c, "", errors.New("JSON 解析失败:"+err.Error()))
			continue
		}
		decryptedMessage, err := decryptMessageWithHMAC(msg1.EncryptedMessage, msg1.Hmac)
		if err != nil {
			sendErr(c, decryptedMessage, errors.New(err.Error()))
			continue
		}
		msg := types.RequestMessage{}
		err = json.Unmarshal([]byte(decryptedMessage), &msg)
		if err != nil {
			sendErr(c, decryptedMessage, errors.New("Request is Error"+err.Error()))
			continue
		}
		if msg.Type == types.MessageTypeHeartbeat { // 维持心跳消息
			// 刷新连接时间
			c.HeartbeatTime = time.Now().Unix()
			continue
		} else {
			msg.AccountUuid = c.AccountId
			logic, err := Logic(msg)
			if err != nil {
				sendErr(c, string(body), errors.New("Error"+err.Error()))
				fmt.Println("Read is err", err)
				continue
			}
			logic.Data = decryptedMessage

			message, err := GetWebSocketMessage(logic)
			if err != nil {
				sendErr(c, "", err)
			}
			SendByClient(c, message)
			continue
		}
	}
}

// ClientManager 客户端管理
type ClientManager struct {
	Clients  map[string]*Client  // 保存连接
	Accounts map[string][]string // 账号和连接关系,map的key是账号id即：AccountId，这里主要考虑到一个账号多个连接
	Mu       *sync.Mutex
	svc      *svc.ServiceContext
}

// Manager 定义一个管理Manager
var Manager = ClientManager{
	Clients:  make(map[string]*Client),  // 参与连接的用户，出于性能的考虑，需要设置最大连接数
	Accounts: make(map[string][]string), // 账号和连接关系
	Mu:       new(sync.Mutex),
}

func Setup(svc *svc.ServiceContext) {
	Manager.svc = svc
	go func() {
		defer func() {
			if r := recover(); r != nil {
				log.Println(r)
			}
		}()
		Heartbeat()
	}()
}

func sendErr(c *Client, s string, err error) {
	replyMsg := types.Response{
		Type: types.SysErr,
		Data: s,
		Body: response.Body{
			Code:    500,
			Message: "Error" + err.Error(),
		},
	}
	message, err := GetWebSocketMessage(&replyMsg)
	if err != nil {
		sendErr(c, "", err)
	}
	SendByClient(c, message)
}
func HandleWebSocket(w http.ResponseWriter, r *http.Request) {
	token := r.URL.Query().Get("user_token")
	if token == "" {

		w.WriteHeader(401)
		response.Response(w, nil, errors.New("invalid user_token"))
		return
	}
	Claims, err := comJwt.BackUserByToken(token, Manager.svc.Config.Auth)
	if err != nil {
		w.WriteHeader(500)
		response.Response(w, nil, errors.New("token is error"))
		return
	}
	conn, err := (&websocket.Upgrader{
		// 1. 解决跨域问题
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}).Upgrade(w, r, nil) // 升级
	if err != nil {
		response.Response(w, nil, err)
		return
	}

	ConnId := xid.New().String()
	client := &Client{
		ID:            ConnId, // 连接id
		AccountId:     Claims.UUID,
		HeartbeatTime: time.Now().Unix(),
		Socket:        conn,
	}
	addClient(client)

}
func addClient(client *Client) {
	Manager.Clients[client.ID] = client
	hm := Manager.Accounts[client.AccountId]
	if hm == nil {
		Manager.Accounts[client.AccountId] = make([]string, 0)
	}
	hm = append(hm, client.ID)
	Manager.Accounts[client.AccountId] = hm
	resp := types.Response{
		Type: types.MessageTypeRegister,
		Body: response.Body{
			Code:    1,
			Data:    client.ID,
			Message: "链接成功",
		},
	}
	fmt.Println("加入成功", client.AccountId)
	message, err := GetWebSocketMessage(&resp)
	if err != nil {
		sendErr(client, "", err)
	}
	SendByClient(client, message)
	client.Read()
}

// Heartbeat 维持心跳
func Heartbeat() {
	for {
		// 获取所有的Clients
		Manager.Mu.Lock()
		clients := make([]*Client, 0)
		for _, a := range Manager.Clients {
			clients = append(clients, a)
		}
		Manager.Mu.Unlock()
		for _, c := range clients {
			if time.Now().Unix()-c.HeartbeatTime > types.HeartbeatTime {
				UnAccountBind(types.UnAccountBindData{
					Id:   c.ID,
					Uuid: c.AccountId,
				})
			} else {
				resp := types.Response{
					Type: types.MessageTypeHeartbeat,
					Body: response.Body{
						Code:    1,
						Message: "发送心跳",
					},
				}
				message, err := GetWebSocketMessage(&resp)
				if err != nil {
					sendErr(c, "", err)
				}
				SendByClient(c, message)
			}
		}

		time.Sleep(time.Second * types.HeartbeatCheckTime)
	}
}
func UnAccountBind(c types.UnAccountBindData) {
	ls := Manager.Clients[c.Id]
	if ls != nil {
		_ = ls.Socket.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "Server is closing"))
		delete(Manager.Clients, ls.ID)

		if len(Manager.Accounts[ls.AccountId]) > 0 {
			for k, clientId := range Manager.Accounts[ls.AccountId] {
				if clientId == ls.ID { // 找到绑定客户端Id
					Manager.Accounts[ls.AccountId] = append(Manager.Accounts[ls.AccountId][:k], Manager.Accounts[ls.AccountId][k+1:]...)
				}
			}
		}
	}

}

// GetClient 根据账号获取连接
func GetClient(accountId string) []*Client {
	clients := make([]*Client, 0)

	Manager.Mu.Lock()
	defer Manager.Mu.Unlock()

	if len(Manager.Accounts[accountId]) > 0 {
		for _, clientId := range Manager.Accounts[accountId] {
			if c, ok := Manager.Clients[clientId]; ok {
				clients = append(clients, c)
			}
		}
	}

	return clients
}

func SendByAccounts(accounts []string, msg []byte) {
	for _, accountId := range accounts {
		// 获取连接id
		clients := GetClient(accountId)
		// 发送消息
		for _, c := range clients {
			SendByClient(c, msg)
		}
	}
}
func sendErrByAccounts(accounts []string, s string, err error) {
	replyMsg := types.Response{
		Type: types.SysErr,
		Data: s,
		Body: response.Body{
			Code:    500,
			Message: "Error" + err.Error(),
		},
	}
	message, err := GetWebSocketMessage(&replyMsg)

	for _, accountId := range accounts {
		// 获取连接id
		clients := GetClient(accountId)
		// 发送消息
		for _, c := range clients {
			SendByClient(c, message)
		}
	}
}
func SendByClient(c *Client, msg []byte) {
	_ = c.Socket.WriteMessage(websocket.TextMessage, msg)
}

func GetWebSocketMessage(message *types.Response) ([]byte, error) {

	msg, err := json.Marshal(message)
	if err != nil {
		return nil, err
	}
	encryptedResponse, hmac1 := encryptAndSign(string(msg))
	responseData := WebSocketMessage{
		EncryptedMessage: encryptedResponse,
		Hmac:             hmac1,
	}
	responseDataBt, err := json.Marshal(responseData)
	if err != nil {
		return nil, err
	}
	return responseDataBt, nil
}
func computeHMAC(message string) string {
	mac := hmac.New(sha256.New, hmacKey)
	mac.Write([]byte(message))
	return base64.StdEncoding.EncodeToString(mac.Sum(nil))
}

func encryptAndSign(message string) (string, string) {
	// AES 加密
	encryptedMessage, err := encryptMessage(message)
	if err != nil {
		log.Fatalf("加密失败: %v", err)
	}

	// 生成 HMAC
	hmac1 := computeHMAC(encryptedMessage)

	return encryptedMessage, hmac1
}

func encryptMessage(message string) (string, error) {
	block, err := aes.NewCipher(secretKey)
	if err != nil {
		return "", fmt.Errorf("创建 AES 密钥失败: %v", err)
	}

	// 生成随机 IV
	iv := make([]byte, aes.BlockSize)
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return "", fmt.Errorf("生成随机 IV 失败: %v", err)
	}

	// 填充数据
	paddedMessage := pad([]byte(message))

	// 加密
	ciphertext := make([]byte, len(paddedMessage))
	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(ciphertext, paddedMessage)
	// 拼接 IV 和密文，并进行 Base64 编码
	finalMessage := append(iv, ciphertext...)
	return base64.StdEncoding.EncodeToString(finalMessage), nil
}

func pad(src []byte) []byte {
	padding := aes.BlockSize - len(src)%aes.BlockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(src, padtext...)
}

// decryptMessageWithHMAC 解密 消息并验证 HMAC
func decryptMessageWithHMAC(encryptedMessage string, receivedHMAC string) (string, error) {
	// 1. 将 Base64 编码的密文解析成字节数组
	encryptedBytes, err := base64.StdEncoding.DecodeString(encryptedMessage)
	if err != nil {
		return "", fmt.Errorf("failed to decode base64 string: %v", err)
	}

	// 2. 提取 IV 和密文
	if len(encryptedBytes) < aes.BlockSize {
		return "", fmt.Errorf("invalid ciphertext length")
	}
	iv := encryptedBytes[:aes.BlockSize]         // 前 16 字节是 IV
	ciphertext := encryptedBytes[aes.BlockSize:] // 剩下的是密文

	// 3. 计算 HMAC
	calculatedHMAC := computeHMAC1(encryptedBytes)

	// 4. 验证 HMAC
	if !hmac.Equal([]byte(calculatedHMAC), []byte(receivedHMAC)) {
		return "", fmt.Errorf("HMAC validation failed: data may have been tampered with")
	}

	// 5. 解密密文
	plaintext, err := decryptAES(ciphertext, iv)
	if err != nil {
		return "", err
	}

	// 6. 返回解密后的消息
	return string(plaintext), nil
}

// 使用 AES 解密
func decryptAES(ciphertext, iv []byte) ([]byte, error) {
	block, err := aes.NewCipher(secretKey)
	if err != nil {
		return nil, fmt.Errorf("failed to create AES cipher: %v", err)
	}

	// CBC 模式解密
	mode := cipher.NewCBCDecrypter(block, iv)
	plaintext := make([]byte, len(ciphertext))
	mode.CryptBlocks(plaintext, ciphertext)

	// 去除填充（PKCS#7 填充）
	plaintext = pkcs7UnPadding(plaintext)
	return plaintext, nil
}

// 使用 PKCS#7 去除填充
func pkcs7UnPadding(plaintext []byte) []byte {
	length := len(plaintext)
	padding := int(plaintext[length-1])
	return plaintext[:length-padding]
}

// 计算 HMAC
func computeHMAC1(data []byte) string {
	mac := hmac.New(sha256.New, hmacKey)
	mac.Write(data)
	return base64.StdEncoding.EncodeToString(mac.Sum(nil)) // Base64 编码的 HMAC
}
